from random import choice
from itertools import chain


from torch import tensor, long as tlong
from torch.utils.data import Dataset

from .graph import Graph

class GraphDataset(Dataset):
    def __init__(self, graph:Graph):
        self.graph = graph

    def __getitem__(self, _):
        pos, neg = self.sample_triple_pair(self.graph)
        return tuple(chain(pos, neg))

    def __len__(self):
        return len(self.graph)

    @staticmethod
    def sample_triple_pair(graph: Graph):
        positive = choice(graph.triples_list)
        
        negative = (positive[0], positive[1], choice(graph.all_entids))
        while negative in graph:
            negative = (positive[0], positive[1], choice(graph.all_entids))

        return positive, negative